package io.gitee.hefren.hhsharding.datasource;

import com.alibaba.druid.pool.DruidDataSourceFactory;
import io.gitee.hefren.hhsharding.config.HhShardingProperties;
import io.gitee.hefren.hhsharding.core.HhShardingContext;
import io.gitee.hefren.hhsharding.core.HhShardingResult;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * @Date 2024/8/16
 * @Author lifei
 */
public class HhShardingDataSource extends AbstractRoutingDataSource {
    private static final Logger log = LoggerFactory.getLogger(HhShardingDataSource.class);

    public HhShardingDataSource(HhShardingProperties hhShardingProperties) {
        Map<Object, Object> targetDataSourceMap = hhShardingProperties.getDatasources().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, item -> {
            try {
                return DruidDataSourceFactory.createDataSource(item.getValue());
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }));

        setTargetDataSources(targetDataSourceMap);
        setDefaultTargetDataSource(targetDataSourceMap.values().iterator().next());
    }

    @Override
    protected Object determineCurrentLookupKey() {
        HhShardingResult hhShardingResult = HhShardingContext.getHhShardingResult();
        log.debug("===> determineCurrentLookupKey: {}", hhShardingResult);
        if (Objects.isNull(hhShardingResult)) {
            return null;
        }
        return hhShardingResult.getTargetDataSourceName();
    }
}
